
package cartagows.wsframework.wssecurity.core;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Set;
import java.util.Vector;
import java.util.Map.Entry;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.ws.security.WSConstants;
import org.apache.ws.security.WSEncryptionPart;
import org.apache.ws.security.WSSecurityEngineResult;
import org.apache.ws.security.WSSecurityException;
import org.apache.ws.security.components.crypto.Crypto;
import org.apache.ws.security.components.crypto.CryptoFactory;
import org.apache.ws.security.handler.WSHandlerResult;
import org.apache.ws.security.message.WSSecEncryptedKey;
import org.apache.ws.security.message.WSSecSignature;
import org.apache.ws.security.message.WSSecSignatureConfirmation;
import org.apache.ws.security.message.WSSecTimestamp;
import org.apache.ws.security.message.WSSecUsernameToken;
import org.apache.ws.security.util.WSSecurityUtil;
import org.w3c.dom.Document;
import org.w3c.dom.Element;

import cartagows.wsframework.wssecurity.policy.Constants;
import cartagows.wsframework.wssecurity.policy.PolicyData;
import cartagows.wsframework.wssecurity.policy.model.SupportingToken;
import cartagows.wsframework.wssecurity.policy.model.Token;
import cartagows.wsframework.wssecurity.policy.model.UsernameToken;
import cartagows.wsframework.wssecurity.policy.model.X509Token;
import cartagows.wsframework.wssecurity.utils.SecurityUtil;

public abstract class BindingBuilder {
	private static Log log = LogFactory.getLog(BindingBuilder.class);

	private Element insertionLocation;
	protected String mainSigId = null;
	protected ArrayList encryptedTokensIdList = new ArrayList();
	protected Element timestampElement;
	protected Element mainRefListElement;


	 protected void addTimestamp(MessageData md) {
		 log.debug("Adding timestamp");

		 WSSecTimestamp timestampBuilder = new WSSecTimestamp();

		 timestampBuilder.setTimeToLive(md.getTimeToLive());

		 // add the Timestamp to the SOAP Enevelope
		 timestampBuilder.build(md.getDocument(), md.getSecHeader());

		 log.debug("Timestamp id: " + timestampBuilder.getId());

		 md.setTimestampId(timestampBuilder.getId());

		 this.timestampElement = timestampBuilder.getElement();
		 log.debug("Adding timestamp: DONE");
	 }

	 protected WSSecUsernameToken addUsernameToken(MessageData md, UsernameToken token) throws SecurityException {

		 log.debug("Adding a UsernameToken");

		 PolicyData pd = md.getPolicyData();

		 //Get the user

		 String user = md.getUtUsername();

		 if(user != null && !"".equals(user)) {
			 log.debug("User : " + user);

			 // If NoPassword property is set we don't need to set the password
			 if (token.isNoPassword()) {
				 WSSecUsernameToken utBuilder = new WSSecUsernameToken();
				 utBuilder.setUserInfo(user, null);
				 utBuilder.setPasswordType(null);

				 return utBuilder;
			 }

			 //Get the password

			 //First check options object for a password
			 String password = md.getUtPassword();

			 if(password != null && !"".equals(password)) {
				 //If the password is available then build the token

				 WSSecUsernameToken utBuilder = new WSSecUsernameToken();
				 if (token.isHashPassword()) {
					 utBuilder.setPasswordType(WSConstants.PASSWORD_DIGEST);  
				 } else {
					 utBuilder.setPasswordType(WSConstants.PASSWORD_TEXT);
				 }

				 utBuilder.setUserInfo(user, password);

				 return utBuilder;
			 } else {
				 //If there's no password then throw an exception
				 throw new SecurityException("noPasswordForUser");
			 }

		 } else {
			 log.debug("No user value specified in the configuration");
			 throw new SecurityException("userMissing");
		 }

	 }

	 protected WSSecEncryptedKey getEncryptedKeyBuilder(MessageData md, Token token) throws SecurityException {

		 PolicyData pd = md.getPolicyData();
		 Document doc = md.getDocument();

		 WSSecEncryptedKey encrKey = new WSSecEncryptedKey();

		 try {
			 SecurityUtil.setKeyIdentifierType(pd, encrKey, token);
			 SecurityUtil.setEncryptionUser(md, encrKey);
			 encrKey.setKeySize(pd.getAlgorithmSuite().getMaximumSymmetricKeyLength());
			 encrKey.setKeyEncAlgo(pd.getAlgorithmSuite().getAsymmetricKeyWrap());
			 Crypto crypto = CryptoFactory.getInstance(md.getCryptoEncProperties());
			 encrKey.prepare(doc, crypto);

			 return encrKey;
		 } catch (WSSecurityException e) {
			 throw new SecurityException("errorCreatingEncryptedKey");
		 }
	 }

	 protected WSSecSignature getSignatureBuilder(MessageData md, Token token) throws SecurityException {

		 PolicyData pd = md.getPolicyData();

		 WSSecSignature sig = new WSSecSignature();
		 checkForX509PkiPath(sig, token);

		 log.debug("Token inclusion: " + token.getInclusion());

		 SecurityUtil.setKeyIdentifierType(pd, sig, token);

		 String user = md.getSignatureUser();
		 String password = md.getSignatureDecryptionPassword();

		 sig.setUserInfo(user, password);
		 sig.setSignatureAlgorithm(pd.getAlgorithmSuite().getAsymmetricSignature());
		 sig.setSigCanonicalization(pd.getAlgorithmSuite().getInclusiveC14n());
		 Crypto crypto = CryptoFactory.getInstance(md.getCryptoDecProperties());
		 try {
			 sig.prepare(md.getDocument(), crypto, md.getSecHeader());
		 } catch (WSSecurityException e) {
			 throw new SecurityException("errorInSignatureWithX509Token");
		 }

		 return sig;
	 }

	 protected HashMap handleSupportingTokens(MessageData md, SupportingToken suppTokens)
	 throws SecurityException {

		 //Create the list to hold the tokens
		 HashMap endSuppTokMap = new HashMap();

		 if(suppTokens != null && suppTokens.getTokens() != null &&
				 suppTokens.getTokens().size() > 0) {
			 log.debug("Processing supporting tokens");

			 ArrayList tokens = suppTokens.getTokens();
			 for (Iterator iter = tokens.iterator(); iter.hasNext();) {
				 Token token = (Token) iter.next();

				 if(token instanceof X509Token) {

					 //We have to use a cert
					 //Prepare X509 signature
					 WSSecSignature sig = this.getSignatureBuilder(md, token);
					 Element bstElem = sig.getBinarySecurityTokenElement();
					 if(bstElem != null) {   
						 bstElem = SecurityUtil.insertSiblingAfter(md, 
								 this.getInsertionLocation(), bstElem);
						 this.setInsertionLocation(bstElem);

						 if (suppTokens.isEncryptedToken()) {
							 this.encryptedTokensIdList.add(sig.getBSTTokenId());
						 }
					 }
					 endSuppTokMap.put(token, sig);

				 } else if(token instanceof UsernameToken) {
					 WSSecUsernameToken utBuilder = addUsernameToken(md, (UsernameToken)token);

					 utBuilder.prepare(md.getDocument());

					 //Add the UT
					 Element elem = utBuilder.getUsernameTokenElement();
					 SecurityUtil.insertSiblingAfter(md, this.getInsertionLocation(), elem);

					 //encryptedTokensIdList.add(utBuilder.getId());

					 //Move the insert location to the next element
					 this.setInsertionLocation(elem);
					 endSuppTokMap.put(token, utBuilder.getId());
					 /*Date now = new Date();
                    try {
                    	org.apache.rahas.Token tempTok = new org.apache.rahas.Token(utBuilder.getId(), (OMElement) elem, now, new Date(now.getTime() + 300000));
                        endSuppTokMap.put(token, tempTok);
                    } catch (TrustException e) {
                        throw new SecurityException("errorCreatingRahasToken");
                    }*/
				 }
			 }
		 }

		 return endSuppTokMap;
	 }

	 protected Vector addSignatureParts(HashMap tokenMap, Vector sigParts) throws SecurityException {

		 Set entrySet = tokenMap.entrySet();

		 for (Iterator iter = entrySet.iterator(); iter.hasNext();) {
			 Object tempTok =  ((Entry)iter.next()).getValue();
			 WSEncryptionPart part = null;

			 if(tempTok instanceof String) {

				 part = new WSEncryptionPart((String) tempTok);

			 } else if(tempTok instanceof WSSecSignature) {
				 WSSecSignature tempSig = (WSSecSignature) tempTok;
				 if(tempSig.getBSTTokenId() != null) {
					 part = new WSEncryptionPart(tempSig.getBSTTokenId());
				 }
			 } else {

				 throw new SecurityException("UnsupportedTokenInSupportingToken");  
			 }
			 sigParts.add(part);
		 }

		 return sigParts;
	 }


	 public Element getInsertionLocation() {
		 return insertionLocation;
	 }

	 public void setInsertionLocation(Element insertionLocation) {
		 this.insertionLocation = insertionLocation;
	 }

	 protected Vector doEndorsedSignatures(MessageData md, HashMap tokenMap) throws SecurityException {
	        
	        Set tokenSet = tokenMap.keySet();
	        
	        Vector sigValues = new Vector();
	        
	        for (Iterator iter = tokenSet.iterator(); iter.hasNext();) {
	            
	            Token token = (Token)iter.next();
	            
	            Object tempTok = tokenMap.get(token);
	            
	            Vector sigParts = new Vector();
	            sigParts.add(new WSEncryptionPart(this.mainSigId));
	            
	            if (tempTok instanceof WSSecSignature) {
	                WSSecSignature sig = (WSSecSignature)tempTok;
	                if(md.getPolicyData().isTokenProtection() &&
	                        sig.getBSTTokenId() != null) {
	                    sigParts.add(new WSEncryptionPart(sig.getBSTTokenId()));
	                }
	                
	                try {
	                    sig.addReferencesToSign(sigParts, md.getSecHeader());
	                    sig.computeSignature();
	                    
	                    this.setInsertionLocation(SecurityUtil.insertSiblingAfter(
	                            md, 
	                            this.getInsertionLocation(), 
	                            sig.getSignatureElement()));
	                    
	                } catch (WSSecurityException e) {
	                    throw new SecurityException("errorInSignatureWithX509Token");
	                }
	                sigValues.add(sig.getSignatureValue());
	            }
	        } 

	        return sigValues;
	            
	    }
	 
	 protected void addSignatureConfirmation(MessageData md, Vector sigParts) {

		 if(!md.getPolicyData().isSignatureConfirmation()) {

			 //If we don't require sig confirmation simply go back :-)
			 return;
		 }

		 Document doc = md.getDocument();

		 Vector results = md.getResults();
		 /*
		  * loop over all results gathered by all handlers in the chain. For each
		  * handler result get the various actions. After that loop we have all
		  * signature results in the signatureActions vector
		  */
		 Vector signatureActions = new Vector();
		 for (int i = 0; i < results.size(); i++) {
			 WSHandlerResult wshResult = (WSHandlerResult) results.get(i);

			 WSSecurityUtil.fetchAllActionResults(wshResult.getResults(),
					 WSConstants.SIGN, signatureActions);
			 WSSecurityUtil.fetchAllActionResults(wshResult.getResults(),
					 WSConstants.ST_SIGNED, signatureActions);
			 WSSecurityUtil.fetchAllActionResults(wshResult.getResults(),
					 WSConstants.UT_SIGN, signatureActions);
		 }

		 // prepare a SignatureConfirmation token
		 WSSecSignatureConfirmation wsc = new WSSecSignatureConfirmation();
		 if (signatureActions.size() > 0) {
			 if (log.isDebugEnabled()) {
				 log.debug("Signature Confirmation: number of Signature results: "
						 + signatureActions.size());
			 }
			 for (int i = 0; i < signatureActions.size(); i++) {
				 WSSecurityEngineResult wsr = (WSSecurityEngineResult) signatureActions
				 .get(i);
				 byte[] sigVal = (byte[]) wsr.get(WSSecurityEngineResult.TAG_SIGNATURE_VALUE);
				 wsc.setSignatureValue(sigVal);
				 wsc.prepare(doc);
				 SecurityUtil.appendChildToSecHeader(md, wsc.getSignatureConfirmationElement());
				 if(sigParts != null) {
					 sigParts.add(new WSEncryptionPart(wsc.getId()));
				 }
			 }
		 } else {
			 //No Sig value
			 wsc.prepare(doc);
			 SecurityUtil.appendChildToSecHeader(md, wsc.getSignatureConfirmationElement());
			 if(sigParts != null) {
				 sigParts.add(new WSEncryptionPart(wsc.getId()));
			 }
		 }
	 }

	 private void checkForX509PkiPath(WSSecSignature sig, Token token){
		 if (token instanceof X509Token) {
			 X509Token x509Token = (X509Token) token;
			 if (x509Token.getTokenVersionAndType().equals(Constants.WSS_X509_PKI_PATH_V1_TOKEN10)
					 || x509Token.getTokenVersionAndType().equals(Constants.WSS_X509_PKI_PATH_V1_TOKEN11)) {
				 sig.setUseSingleCertificate(false);
			 }
		 }
	 }


}
